"""
数据预处理模块
实现书法图像的标准化处理，包括灰度化、二值化、骨架提取等
"""
import cv2
import numpy as np
from skimage.morphology import skeletonize, remove_small_objects
from skimage import measure

class ImagePreprocessor:
    def __init__(self):
        self.original_image = None
        self.gray_image = None
        self.binary_image = None
        self.skeleton_image = None
        self.processed_image = None
        
    def load_image(self, image_path):
        """加载图像文件"""
        self.original_image = cv2.imread(image_path)
        if self.original_image is None:
            raise ValueError("无法加载图像，请检查文件路径")
        return self.original_image
    
    def load_image_from_array(self, image_array):
        """从数组中加载图像"""
        self.original_image = image_array
        return self.original_image
    
    def to_grayscale(self):
        """图像灰度化处理"""
        if self.original_image is None:
            raise ValueError("请先加载图像")
        
        # 如果是彩色图像，转换为灰度图
        if len(self.original_image.shape) == 3:
            self.gray_image = cv2.cvtColor(self.original_image, cv2.COLOR_BGR2GRAY)
        else:
            self.gray_image = self.original_image.copy()
        
        return self.gray_image
    
    def to_binary(self, threshold_method=cv2.THRESH_OTSU):
        """
        图像二值化处理
        使用自适应阈值或OTSU方法
        """
        if self.gray_image is None:
            self.to_grayscale()
        
        # 使用OTSU自动阈值法或指定阈值方法
        _, self.binary_image = cv2.threshold(self.gray_image, 0, 255, threshold_method)
        
        # 确保书法为黑色（前景），背景为白色
        if np.sum(self.binary_image == 0) > np.sum(self.binary_image == 255):
            self.binary_image = 255 - self.binary_image
            
        return self.binary_image
    
    def extract_skeleton(self):
        """
        提取骨架
        使用改进的Zhang-Suen算法
        """
        if self.binary_image is None:
            self.to_binary()
        
        # 转换为skimage需要的格式（二值图像，True表示前景）
        binary = self.binary_image == 0
        
        # 应用骨架化算法
        skeleton = skeletonize(binary)
        
        # 转回OpenCV格式
        self.skeleton_image = np.where(skeleton, 0, 255).astype(np.uint8)
        
        return self.skeleton_image
    
    def remove_noise(self, min_size=30):
        """
        去除噪点
        使用连通域分析移除小区域
        """
        if self.binary_image is None:
            self.to_binary()
        
        # 转换为skimage需要的格式
        binary = self.binary_image == 0
        
        # 去除小连通域
        cleaned = remove_small_objects(binary, min_size=min_size)
        
        # 转回OpenCV格式
        cleaned_image = np.where(cleaned, 0, 255).astype(np.uint8)
        
        self.processed_image = cleaned_image
        return self.processed_image
    
    def normalize_size(self, target_height=512, preserve_aspect=True):
        """
        归一化图像大小，保持原始比例
        """
        if self.processed_image is None:
            if self.binary_image is not None:
                self.processed_image = self.binary_image
            else:
                self.to_binary()
                self.processed_image = self.binary_image
        
        h, w = self.processed_image.shape
        
        if preserve_aspect:
            # 保持原始比例
            aspect = w / h
            new_height = target_height
            new_width = int(aspect * new_height)
        else:
            new_height = target_height
            new_width = target_height
        
        # 使用非线性插值算法进行缩放
        self.processed_image = cv2.resize(
            self.processed_image, 
            (new_width, new_height), 
            interpolation=cv2.INTER_LANCZOS4
        )
        
        return self.processed_image
    
    def process_image(self, image, target_height=512):
        """
        完整的图像处理流程
        """
        self.load_image_from_array(image)
        self.to_grayscale()
        self.to_binary()
        self.remove_noise()
        self.extract_skeleton()
        self.normalize_size(target_height)
        
        return {
            'original': self.original_image,
            'grayscale': self.gray_image,
            'binary': self.binary_image,
            'skeleton': self.skeleton_image,
            'processed': self.processed_image
        } 